#include "iam_matop.h"
#include <string.h>

extern void xzsvdc(IAM_RasterView_F64 A, IAM_RasterData_F64 U, IAM_RasterData_F64 S, IAM_RasterData_F64 V);


void IAM_MatOp_transpose_F64(const IAM_RasterView_F64 src_mat, IAM_RasterView_F64 dst_mat)
{
    size_t y, x;
    if (src_mat->m_pdata == dst_mat->m_pdata) {
        size_t N = src_mat->m_size_x * src_mat->m_size_y;
        double* p_buffer = malloc(sizeof(double) * N);
        memcpy(p_buffer, src_mat->m_pdata, sizeof(double) * N);
        
        for (y = 0; y < src_mat->m_size_y; ++y) {
            for (x = 0; x < src_mat->m_size_x; ++x) {
                size_t dst_ind = y + x * src_mat->m_size_y;
                size_t src_ind = x + y * src_mat->m_size_x;
                src_mat->m_pdata[dst_ind] = p_buffer[src_ind];
            }
        }
        free(p_buffer);
        size_t tmp = src_mat->m_size_x;
        src_mat->m_size_x = src_mat->m_size_y;
        src_mat->m_size_y = tmp;
    }
    else {
        size_t N = src_mat->m_size_x * src_mat->m_size_y;
        if (N != dst_mat->m_size_x * dst_mat->m_size_y) {
            fprintf(stderr, "IAM_MatOp_transpose_F64 Fatal Error: invalid size!\n");
            exit(-1);
        }

        for (y = 0; y < src_mat->m_size_y; ++y) {
            for (x = 0; x < src_mat->m_size_x; ++x) {
                size_t dst_ind = y + x * src_mat->m_size_y;
                size_t src_ind = x + y * src_mat->m_size_x;
                dst_mat->m_pdata[dst_ind] = src_mat->m_pdata[src_ind];
            }
        }

        dst_mat->m_size_x = src_mat->m_size_y;
        dst_mat->m_size_y = src_mat->m_size_x;;
    }
}

void IAM_MatOp_svd_F64(const IAM_RasterView_F64 A, IAM_RasterData_F64 U, IAM_RasterData_F64 S, IAM_RasterData_F64 V)
{
    // 拷贝矩阵A
    IAM_RasterData_F64 clone_A = IAM_RasterData_init_F64(A->m_size_x, A->m_size_y);
    memcpy(clone_A->m_p_raw_data, A->m_pdata, sizeof(double)*A->m_size_x*A->m_size_y);

    // FIXME matlab按照列存储, 得交换下的 U 和 V, 矩阵A会被修改(所以我们需要拷贝一份)
    // A = USV', A' = VSU'
    xzsvdc(clone_A->m_view, V, S, U); 
    
    // 输出按列存储, 转置 U 和 V 变为按行存储
    IAM_MatOp_transpose_F64(U->m_view, U->m_view);
    IAM_MatOp_transpose_F64(V->m_view, V->m_view);

    IAM_RasterData_release_F64(&clone_A);
}

void IAM_MatOp_lsqsvd_F64(const IAM_RasterView_F64 U, const IAM_RasterView_F64 S, const IAM_RasterView_F64 V, const IAM_RasterView_F64 b, IAM_RasterView_F64 X)
{
    int M = (int)U->m_size_y;
    int N = (int)V->m_size_y;
    int K = (int)U->m_size_x;

    if (b->m_size_x != 1) {
        fprintf(stderr, "IAM_MatOp_lsqsvd_F64 Fatal Error: b is not a column vector!\n");
        exit(-1);
    }
    if (X->m_size_x != 1) {
        fprintf(stderr, "IAM_MatOp_lsqsvd_F64 Fatal Error: X is not a column vector!\n");
        exit(-1);
    }
    if (b->m_size_y != M) {
        fprintf(stderr, "IAM_MatOp_lsqsvd_F64 Fatal Error: b invalid size !\n");
        exit(-1);
    }
    if (X->m_size_y != N) {
        fprintf(stderr, "IAM_MatOp_lsqsvd_F64 Fatal Error: X invalid size !\n");
        exit(-1);
    }

    // (row,col)
    // U: (M,k) S:(k,k) V:(N,k), k=min(M,N)
    // b: (M,1)
    
    int i, j;
    // beta=U'*b : (k,1)
    IAM_RasterData_F64 beta = IAM_RasterData_init_F64(1, N);
    for (i = 0; i < N; ++i) {
        double val = 0;
        for (j = 0; j < M; ++j) {
            double tmp = IAM_RasterView_getValue_F64(U, i, j);
            val += (b->m_pdata[j] * tmp);
        }
        beta->m_p_raw_data[i] = val;
    }

    // xi = beta / sig
    // beta <== xi : (k,1)
    for (i = 0; i < K; ++i) {
        if (S->m_pdata[i] != 0) {
            beta->m_p_raw_data[i] /= S->m_pdata[i];
        }
        else {
            beta->m_p_raw_data[i] = 0;
        }
    }

    // beta <== xi : (k,1)
    // V:(N,k), k=min(M,N)
    // x = sum( xi'.*V, 2 );
    for (i = 0; i < N; ++i) {
        double val = 0;
        for (j = 0; j < K; ++j) {
            double tmp = IAM_RasterView_getValue_F64(V, j, i);
            val += tmp * beta->m_p_raw_data[j];
        }
        X->m_pdata[i] = val;
    }

    IAM_RasterData_release_F64(&beta);
}

void IAM_MatOp_lsq_F64(const IAM_RasterView_F64 A, const IAM_RasterView_F64 b, IAM_RasterView_F64 X)
{
    IAM_RasterData_F64 U = IAM_RasterData_init0_F64();
    IAM_RasterData_F64 S = IAM_RasterData_init0_F64();
    IAM_RasterData_F64 V = IAM_RasterData_init0_F64();

    IAM_MatOp_svd_F64(A, U, S, V);
    IAM_MatOp_lsqsvd_F64(U->m_view, S->m_view, V->m_view, b, X);

    IAM_RasterData_release_F64(&U);
    IAM_RasterData_release_F64(&S);
    IAM_RasterData_release_F64(&V);
}


